import numpy as np
import onnxruntime as ort
import cv2
import changeONNX
import genpromodel


def infer():

    # 加载ONNX模型
    session = ort.InferenceSession("det.sim.normal.onnx")

    # 准备输入数据
    input_name = session.get_inputs()[0].name  # 获取输入节点的名称

    input_data = cv2.imread("141.jpg") # 示例输入数据
    input_data = cv2.resize(input_data, (960, 960)).astype("float32")
    input_data = input_data.transpose((2, 0, 1))[None]


    # 执行推理
    output_name = session.get_outputs()[0].name  # 获取输出节点的名称
    outputs = session.run([output_name], {input_name: input_data})

    # 处理输出结果
    output_data = np.transpose(outputs[0][0], (1, 2, 0))*255.0  # 获取第一个输出结果
    cv2.imshow("  ", output_data)
    cv2.waitKey(0)
    print(output_data.shape)



if __name__ == "__main__":

    genpromodel.genProModel()
    changeONNX.mergeOnnx()

    infer()